import re
from typing import Any
from preprocess.file_paser.paser_utils import re_filter
from langchain.text_splitter import CharacterTextSplitter

EXTRA_COL_NAME = 'meta_index'


def table_split(split_len, data_text, head_text):
    """
    这里切的是表格的数据部分（data_text），但要保证切割后的数据+表头的长度小于split_len
    :param split_len: 最大切割长度
    :param data_text: 表格的数据部分
    :param head_text: 表头字符串
    :return:
    """
    head_len = len(head_text)
    c_splitter = CharacterTextSplitter(
        separator='\n',
        chunk_size=split_len - head_len,
        chunk_overlap=0
    )
    results = []
    pattern = re.compile('[-| \n]+')
    for sub_table in c_splitter.split_text(data_text):
        txt = re_filter(head_text + sub_table)
        if len(re.sub(pattern, '', txt).strip()) > 0:
            results.append(txt)
    return results


def find_split_indices(table, max_len):
    columns = table.columns.tolist()
    # 检测表头的长度

    indices = []
    # if (total_head_len < max_len):
    #     return None
    s = 0
    e = -1
    # markdown形式所需添加的字符数量 开头| 与 结尾 \n
    col_len = 4
    row_len = [2 for _ in range(table.shape[0])]

    # 补一条index col如果有需求要纵向切割
    extraColLen = len(EXTRA_COL_NAME) + 7
    extraRowLen = len(str(table.shape[0])) + 1

    for i, col in enumerate(columns):
        cur_col_len = len(str(col).replace('\n', '').strip()) + 7

        cur_row_len = []
        for cell in table.iloc[:, i]:
            cell_len = len(str(cell).replace('\n', '').strip()) + 1
            cur_row_len.append(cell_len)

        # print(f'cur_row_len: {cur_row_len}')
        sum_row_len = list(map(lambda x, y: x + y, cur_row_len, row_len))
        # print(f'row_size at {i}:{row_len}')
        if (col_len + cur_col_len + max(sum_row_len) + extraColLen + extraRowLen < max_len):
            e = i
            col_len += cur_col_len
            row_len = sum_row_len
        elif (e >= s):
            # print(f'num_col:{e}')
            # print(f'column_len:{col_len}')
            # print(f'current chunk_size:{col_len + max(row_len)}')
            indices.append((s, e))
            col_len = cur_col_len + 4
            row_len = [2 + item for item in cur_row_len]
            s = i
            e = i
    indices.append((s, len(columns) - 1))
    return indices


def vert_split(indices, table):
    """_summary_
        纵向切割表格
    Args:
        indices (list[tuples]): 很多组切割的index
        table (df): 储存为dataframe的excel表格

    Returns:
        list[dataframes]: 分割后的一系列dataframe
    """
    # indices list只有一个tuple说明表格不需要纵向分割
    if (len(indices) == 1):
        return [table]
    split_table_ls = []
    num_rows = table.shape[0]
    # 分割表格
    for start, end in indices:
        split_df = table.iloc[:, start:end + 1].copy()
        # 添加index列用便于后续合并
        index_col = [i for i in range(num_rows)]
        split_df.insert(0, EXTRA_COL_NAME, index_col)
        split_table_ls.append(split_df)
    return split_table_ls


def process_excel_table(table, max_len):
    """_summary_

    Args:
        table (df): 表格
        max_len (int): 将表格切分为子表格，每个子表格的字符串长度不超过 max_len

    Returns:
        list[df]: markdown字符串形式输出table列表
    """
    table_data = []  # 表数据
    extra_data = []  # 溢出数据
    columns = table.columns.tolist()

    # 遍历table中的一行，生成markdown形式的text
    def traverse_header(columns):
        row_txt = ['|']
        for cid, content in enumerate(columns):
            cell_txt = str(content).replace('\n', '').strip()
            # 单元格为空时用“无”填充
            row_txt.append(cell_txt if len(cell_txt) != 0 else '无')
            row_txt.append('|')
        return row_txt

    # 制作表头(excel表头唯一性可确定)
    head_row_text = traverse_header(columns)
    head_row_text = ''.join(head_row_text).strip()
    head_split = '|' + '-|' * len(columns)
    head_text = head_row_text + '\n' + head_split + '\n'  # 融合表头
    # 制作表身
    for i, row in table.iterrows():
        row_text = ['|']
        for cell_id, cell in enumerate(row):
            cell_text = str(cell).replace('\n', '').strip()
            if len(cell_text) > max_len - len(head_text):
                head = str(columns[cell_id]).replace('\n', '').strip()
                extra_data.append(re_filter(f'{head}: {cell_text}'))
                row_text.append(' |')
            else:
                # 单元格为空时用“无”填充
                row_text.append(cell_text if cell_text != 'nan' else '无')
                row_text.append('|')
        table_data.append(''.join(row_text).strip())

    data_text = '\n'.join(table_data)  # 表格中的数据部分
    table_text = head_text + data_text
    if len(table_text) < max_len:
        return [table_text], extra_data
    # print(head_text)
    sub_tables = table_split(max_len, data_text, head_text)  # 切割 markdown形式的表格字符串
    return sub_tables, extra_data


def prepare_excel_table(table, max_len) -> tuple[list[Any], list[Any]]:
    """
    :param table: df类型
    :param max_len: 如果超过该值则将表格划分为多个子表
    :return: table_ls：子表的markdown字符串，extra_data_ls：单个单元格内容字符串超过max_len的内容列表
    """
    table = table.rename(columns=lambda x: x.replace('Unnamed:', '') if 'Unnamed:' in x else x)
    indices = find_split_indices(table, max_len)
    split_table_ls = vert_split(indices, table)
    table_ls = []
    extra_data_ls = []
    for smaller_table in split_table_ls:
        if smaller_table.empty:
            continue
        sub_tables, extra_data = process_excel_table(smaller_table, max_len)
        table_ls.extend(sub_tables)
        extra_data_ls.extend(extra_data)
    return table_ls, extra_data_ls
